# -*- coding: utf-8 -*-
# created on 2017/01/09

from mathsolver.functions.base import *
from mathsolver.functions.shulie.Fundamental import *
from mathsolver.functions.shulie.QianNXiangHe import ShuLieQiuHe
from sympy.abc import n


# 等差数列型求和
class GetNumSums001(BaseFunction):
    def solver(self, *args):
        sl = args[0]
        target_poly = args[1]
        num = args[2]
        sl_value = sl.value
        if sl_value != 'DengCha':
            raise Exception('Try Error!')
        sl.get_DiTuiGongShi()
        if sl.DiTuiGongShi_Eqs:
            dituigongshi = sl.DiTuiGongShi_Eqs
            if len(dituigongshi) != 1:
                raise Exception('Try Error!')
            eq1 = [dituigongshi[0][0].subs({n: 1}), dituigongshi[0][1].subs({n: 1})]
            eq2 = [dituigongshi[0][0].subs({n: 2}), dituigongshi[0][1].subs({n: 2})]
            new_eqs = [eq1, eq2]
            sl.Eqs_Normalized(BaseEqs(new_eqs))
        else:
            sl.Eqs_Normalized([])
        self.steps += sl.steps_EqsNormalized
        sl.get_JiBenLiang()
        self.steps += sl.steps_JiBenLiang
        self.label.update(sl.label_JiBenLiang)
        poly_name = sl_shulie_name(target_poly)
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(target_poly))
        m = list(set(m))
        if len(poly_name) == 2:
            print(1)
        elif len(poly_name) == 1:
            if str(poly_name[0]) == 'S_':
                print(1)
            else:
                sl.get_TongXiangGongShi()
                self.steps += sl.steps_TongXiang
                tongxiangs = sl.an
                target_values = []
                for ans in tongxiangs:
                    new_target_poly = target_poly
                    for item in m:
                        new_item = sympify(item)
                        name_item = new_item.func
                        value_item = 0
                        if str(name_item) == 'a_':
                            index_item = new_item.args[0]
                            value_item = ans.subs({n: index_item})
                        new_target_poly = new_target_poly.subs({new_item: value_item})
                    target_values.append(new_target_poly)
                self.steps.append(["∴", " 或 ".join(["%s = %s" % (new_latex(target_poly), new_latex(v)) for v in target_values])])
                target_sums = []
                for val in target_values:
                    name = sl_shulie_name(val)
                    if name:
                        stepsolver = ShuLieQiuHe().solver([name[0](n), val], num)
                    else:
                        self.steps.append(["", "设%s = %s, 则" % (new_latex(sympify('b_(n)')), new_latex(val))])
                        stepsolver = ShuLieQiuHe().solver([sympify('b_(n)'), val], num)
                    self.steps += stepsolver.steps
                    self.label.update(stepsolver.label)
                    val_sums = stepsolver.output[0].sympify()
                    for sval in val_sums:
                        target_sums.append(sval)
                self.output.append(BaseValues(target_sums))
        return self


# 等比数列型求和
class GetNumSums002(BaseFunction):
    def solver(self, *args):
        sl = args[0]
        target_poly = args[1]
        num = args[2]
        sl_value = sl.value
        if sl_value != 'DengBi':
            raise Exception('Try Error!')
        sl.get_DiTuiGongShi()
        if sl.DiTuiGongShi_Eqs:
            dituigongshi = sl.DiTuiGongShi_Eqs
            if len(dituigongshi) != 1:
                raise Exception('Try Error!')
            eq1 = [dituigongshi[0][0].subs({n: 1}), dituigongshi[0][1].subs({n: 1})]
            eq2 = [dituigongshi[0][0].subs({n: 2}), dituigongshi[0][1].subs({n: 2})]
            new_eqs = [eq1, eq2]
            sl.Eqs_Normalized(BaseEqs(new_eqs))
        else:
            sl.Eqs_Normalized([])
        self.steps += sl.steps_EqsNormalized
        sl.get_JiBenLiang()
        self.steps += sl.steps_JiBenLiang
        poly_name = sl_shulie_name(target_poly)
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(target_poly))
        m = list(set(m))
        if len(poly_name) == 2:
            raise Exception("to do")
        elif len(poly_name) == 1:
            if str(poly_name[0]) == 'S_':
                raise Exception("to do")
            else:
                sl.get_TongXiangGongShi()
                self.steps += sl.steps_TongXiang
                tongxiangs = sl.an
                target_values = []
                for ans in tongxiangs:
                    new_target_poly = target_poly
                    for item in m:
                        new_item = sympify(item)
                        name_item = new_item.func
                        value_item = 0
                        if str(name_item) == 'a_':
                            index_item = new_item.args[0]
                            value_item = ans.subs({n: index_item})
                        new_target_poly = new_target_poly.subs({new_item: value_item})
                    target_values.append(new_target_poly)
                self.steps.append(["∴", " 或 ".join(["%s = %s" % (new_latex(target_poly), new_latex(v)) for v in target_values])])
                target_sums = []
                for val in target_values:
                    self.steps.append(["", "设%s = %s, 则" % (new_latex(sympify('b_(n)')), new_latex(val))])
                    stepsolver = ShuLieQiuHe().solver([sympify('b_(n)'), val], 1, num)
                    self.steps += stepsolver.steps
                    self.label.update(stepsolver.label)
                    val_sums = stepsolver.output[0].sympify()
                    for sval in val_sums:
                        target_sums.append(sval)
                self.output.append(BaseValues(target_sums))
        return self


class GetNumSums003(BaseFunction):
    """
    数列{a_{n}}满足a_{n+1}+(-1)^{n}a_{n}=2n-1,则{a_{n}}的前60项和为()
    """
    def solver(self, *args):
        sl = args[0]
        sl.get_name()
        target_poly = args[1]
        num = args[2]
        if sl.type != 'ShuLie':
            raise Exception('Try Error!')
        assert self.known['Sequations']
        seqs = self.known['Sequations']
        if self.known['inequations']:
            for ineq in self.known['inequations']:
                left, op, right = ineq
                new_ineq = [sympify(left), op, sympify(right)]
                if new_ineq not in seqs:
                    seqs.append(new_ineq)
        if self.known['equations']:
            for eq in self.known['equations']:
                if eq not in seqs:
                    seqs.append(eq)
        sl.eqs_update(seqs)
        sl.get_shouxiang()
        sl.get_dituigongshi()
        sl.get_xiangs()
        poly_name = sl_shulie_name(target_poly)
        p = r'[a-zA-Z]_\(.*?\)'
        m = re.findall(p, str(target_poly))
        m = list(set(m))
        if len(poly_name) == 2:
            raise Exception('Try Error!')
        elif len(poly_name) == 1:
            if str(poly_name[0]) == 'S_':
                raise Exception('Try Error!')
            else:
                sl.get_tongxianggongshi()
                self.steps += sl.steps_TongXiang
                self.label.update(sl.label_TongXiang)
                tongxiangs = sl.an
                target_values = []
                for ans in tongxiangs:
                    new_target_poly = target_poly
                    for item in m:
                        new_item = sympify(item)
                        name_item = new_item.func
                        value_item = 0
                        if str(name_item) == 'a_':
                            index_item = new_item.args[0]
                            value_item = ans.subs({n: index_item})
                        new_target_poly = new_target_poly.subs({new_item: value_item})
                    target_values.append(new_target_poly)
                self.steps.append(["∴", " 或 ".join(["%s = %s" % (new_latex(target_poly), new_latex(v)) for v in target_values])])
                target_sums = []
                for val in target_values:
                    self.steps.append(["", "设%s = %s, 则" % (new_latex(sympify('b_(n)')), new_latex(val))])
                    stepsolver = ShuLieQiuHe().solver([sympify('b_(n)'), val], num)
                    self.steps += stepsolver.steps
                    self.label.update(stepsolver.label)
                    val_sums = stepsolver.output[0].sympify()
                    for sval in val_sums:
                        target_sums.append(sval)
                self.output.append(BaseValues(target_sums))
        elif not target_poly:
            sl.get_TongXiangGongShi()
            self.steps += sl.steps_TongXiang
            self.label.update(sl.label_TongXiang)
            tongxiangs = sl.an
            target_sums = []
            for ans in tongxiangs:
                stepsolver = ShuLieQiuHe().solver([sympify('a_(n)'), ans], num)
                self.steps += stepsolver.steps
                self.label.update(stepsolver.label)
                val_sums = stepsolver.output[0].sympify()
                for sval in val_sums:
                    target_sums.append(sval)
            self.output.append(BaseValues(target_sums))
        return self


# 求数列前n项和
class ShuLieGetNumSums(BaseFunction):
    CLS = [GetNumSums001, GetNumSums002, GetNumSums003]

    def solver(self, *args):
        sl = self.search(str('a_'))
        poly = None
        num = None
        if len(args) == 2:
            poly = args[0].sympify()
            num = args[1].sympify()
        elif len(args) == 1:
            eqs = sl.Eqs
            name = sympify('a_')
            for eq in eqs:
                expr = eq[0] - eq[1]
                name = sl_shulie_name(expr)
            poly = name[0](n)
            num = args[0].sympify()
        r = None
        for cl in ShuLieGetNumSums.CLS:
            try:
                r = cl(known=self.known, verbose=True).solver(sl, poly, num)
                break
            except Exception:
                pass
        if not r:
            raise 'try fail'
        return r
